package io.github.javpower.knowledgerag.core;

import cn.hutool.core.collection.ListUtil;
import io.github.javpower.vectorexclient.util.GsonUtil;
import okhttp3.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Vectorizer {
    private static final String OLLAMA_API_URL = "http://localhost:11434/api/embed"; // 修改为你的 Ollama API URL

    public static List<List<Float>> vectorizeChunks(String OLLAMA_API_URL,String model,List<String> chunks) {
        List<List<Float>> vectors = new ArrayList<>();
        for (String chunk : chunks) {
            List<Float> vector = getVector(OLLAMA_API_URL,model,chunk);
            vectors.add(vector);
        }
        return vectors;
    }

    public static List<Float> getVector(String OLLAMA_API_URL,String model,String text) {
        OkHttpClient client = new OkHttpClient();
        MediaType mediaType = MediaType.parse("application/json");
        Map<String, Object> map=new HashMap<>();
        map.put("model",model);
        map.put("input", ListUtil.of(text));
        String json = GsonUtil.toJson(map);
        RequestBody body = RequestBody.create(json, mediaType);
        Request request = new Request.Builder()
                .url(OLLAMA_API_URL+"/embed")
                .post(body)
                .addHeader("Content-Type", "application/json")
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected code " + response);
            }
            String responseBody = response.body().string();
            return parseVectorFromJson(responseBody);
        } catch (IOException e) {
            e.printStackTrace();
            return null;
        }
    }

    private static List<Float> parseVectorFromJson(String json) {
        try {
            Map map = GsonUtil.fromJson(json, Map.class);
            List<List<Float>> embeddings = (List<List<Float>>)map.get("embeddings");
            List<Float> vector = embeddings.get(0);
            return vector;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }
}